from pydantic import BaseModel,Field
from typing import Optional,Union

from langchain_openai import ChatOpenAI


class UserInfo(BaseModel):
    name:str=Field(description="用户名称")
    age:Optional[int]=Field(description="用户年龄")
    email:str=Field(description="用户邮箱")
    phone:Optional[str]=Field(description="用户手机号")

class ConversationalResponse(BaseModel):
    response: str = Field(description="A conversational response to the user's query")


class FinalResponse(BaseModel):
    final_output: Union[UserInfo,ConversationalResponse]


api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-4o",api_key = api_key,base_url=api_base,temperature=0)

structured_llm = llm.with_structured_output(FinalResponse)

# llm_resp = structured_llm.invoke("介绍下你自己")
# print(llm_resp)
#
# extracted_user_info = structured_llm.invoke("我叫奥特曼，今年38岁，邮箱地址是aoteman#qq.com，电话是1211111111")
# print(extracted_user_info)